#include <stdio.h>
#include <MNN/ImageProcess.hpp>
#include <MNN/expr/Module.hpp>
#include <MNN/expr/Executor.hpp>
#include <MNN/expr/ExprCreator.hpp>
#include <MNN/expr/Executor.hpp>
#include <ctime>
#include <cv/cv.hpp>

using namespace MNN;
using namespace MNN::Express;
using namespace MNN::CV;

int main(int argc, const char* argv[]) {
    std::srand(std::time(0)); // 使用当前时间作为随机数种子
    if (argc < 4) {
        MNN_PRINT("Usage: ./yolov5_demo.out face.mnn keypoint.mnn input.jpg [forwardType] [precision] [thread]\n");
        return 0;
    }
    int thread = 4;
    int precision = 0;
    int forwardType = MNN_FORWARD_CPU;
    if (argc >= 5) {
        forwardType = atoi(argv[4]);
    }
    if (argc >= 6) {
        precision = atoi(argv[5]);
    }
    if (argc >= 7) {
        thread = atoi(argv[6]);
    }
    MNN::ScheduleConfig sConfig;
    sConfig.type = static_cast<MNNForwardType>(forwardType);
    sConfig.numThread = thread;
    BackendConfig bConfig;
    bConfig.precision = static_cast<BackendConfig::PrecisionMode>(precision);
    sConfig.backendConfig = &bConfig;
    std::shared_ptr<Executor::RuntimeManager> rtmgr = std::shared_ptr<Executor::RuntimeManager>(Executor::RuntimeManager::createRuntimeManager(sConfig));
    if(rtmgr == nullptr) {
        MNN_ERROR("Empty RuntimeManger\n");
        return 0;
    }
    rtmgr->setCache(".cachefile");
    std::shared_ptr<Executor::RuntimeManager> rtmgr1 = std::shared_ptr<Executor::RuntimeManager>(Executor::RuntimeManager::createRuntimeManager(sConfig));
    if(rtmgr1 == nullptr) {
        MNN_ERROR("Empty RuntimeManger\n");
        return 0;
    }
    rtmgr1->setCache(".cachefile");
    std::shared_ptr<Module> net_face(Module::load(std::vector<std::string>{}, std::vector<std::string>{}, argv[1], rtmgr));
    std::shared_ptr<Module> net_keypoint(Module::load(std::vector<std::string>{}, std::vector<std::string>{}, argv[2], rtmgr1));
    auto original_image = imread(argv[3]);//获取原始图像
    auto dims = original_image->getInfo()->dim;
    int oih = dims[0];
    int oiw = dims[1];
    int och = dims[2];
    printf("image shape: %d x %d x %d\n", oih, oiw, och);
    int len = oih > oiw ? oih : oiw;
    float scale = len / 640.0;
    std::vector<int> padvals { 0, len - oih, 0, len - oiw, 0, 0 };//左、上、右、下、前、后
    auto pads = _Const(static_cast<void*>(padvals.data()), {3, 2}, NCHW, halide_type_of<int>());
    auto image = _Pad(original_image, pads, CONSTANT);
    image = resize(image, Size(640, 640), 0, 0, INTER_LINEAR, -1, {0., 0., 0.}, {1./255., 1./255., 1./255.});
    auto input = _Unsqueeze(image, {0});
    input = _Convert(input, NC4HW4);
    auto outputs = net_face->onForward({input});
    auto output = _Convert(outputs[0], NCHW);
    output = _Squeeze(output);
    // output shape: [25200, 85]; 85 means: [cx, cy, w, h, box_conf, prob * 80]
    // get box_conf > 0.3 output
    printf("dims = "); for (auto d : output->getInfo()->dim) printf("%d, ", d); printf("\n");
    auto box_conf = _GatherV2(output, _Scalar<int>(4), _Scalar<int>(1));
    auto has_object = _Greater(box_conf, _Scalar<float>(0.1));
    auto idx = Express::_Where(has_object);
    output = Express::_GatherND(output, idx);
    printf("dims = "); for (auto d : output->getInfo()->dim) printf("%d, ", d); printf("\n");
    output = _Transpose(output, {1, 0}); // to [85, 25200]
    auto cx = _Gather(output, _Scalar<int>(0));
    auto cy = _Gather(output, _Scalar<int>(1));
    auto w = _Gather(output, _Scalar<int>(2));
    auto h = _Gather(output, _Scalar<int>(3));
    auto _box_conf = _Gather(output, _Scalar<int>(4));
    std::vector<int> startvals { 5, 0 };
    auto start = _Const(static_cast<void*>(startvals.data()), {2}, NCHW, halide_type_of<int>());
    std::vector<int> sizevals { -1, -1 };
    auto size = _Const(static_cast<void*>(sizevals.data()), {2}, NCHW, halide_type_of<int>());
    auto probs = _Slice(output, start, size);
    // [cx, cy, w, h] -> [y0, x0, y1, x1]
    auto x0 = cx - w * _Const(0.5);
    auto y0 = cy - h * _Const(0.5);
    auto x1 = cx + w * _Const(0.5);
    auto y1 = cy + h * _Const(0.5);
    auto boxes = _Stack({x0, y0, x1, y1}, 1);
    auto scores = _ReduceMax(probs, {0});
    auto ids = _ArgMax(probs, 0);
    auto result_ids = _Nms(boxes, scores, 100, 0.45, 0.25);
    auto result_ptr = result_ids->readMap<int>();
    auto box_ptr = boxes->readMap<float>();
    auto box_conf_ptr = box_conf->readMap<float>();
    auto ids_ptr = ids->readMap<int>();
    auto score_ptr = scores->readMap<float>();
    for (int i = 0; i < 100; i++) {
        auto idx = result_ptr[i];
        if (idx < 0) break;
        auto x0 = box_ptr[idx * 4 + 0] * scale;
        auto y0 = box_ptr[idx * 4 + 1] * scale;
        auto x1 = box_ptr[idx * 4 + 2] * scale;
        auto y1 = box_ptr[idx * 4 + 3] * scale;
        // auto class_idx = ids_ptr[idx];
        // auto score = score_ptr[idx];

        printf("### box: {%f, %f, %f, %f}\n", x0, y0, x1, y1);
        int x0_int = static_cast<int>(x0);
        int y0_int = static_cast<int>(y0);
        int x1_int = static_cast<int>(x1);
        int y1_int = static_cast<int>(y1);
        //获取框的大小
        int bbox_width = x1_int - x0_int;
        int bbox_height = y1_int - y0_int;
        auto face_width = (1 + 2 * 0.2) * bbox_width;//将两边增加0.2
        std::vector<int> center {static_cast<int>((x0 + x1) / 2), static_cast<int>((y0 + y1) / 2)};
        int x0_new = static_cast<int>(center[0] - face_width / 2);
        int y0_new = static_cast<int>(center[1] - face_width / 2);
        int x1_new = static_cast<int>(center[0] + face_width / 2);
        int y1_new = static_cast<int>(center[1] + face_width / 2);
        // 定义裁剪区域
        std::vector<int> start = {y0_new, x0_new, 0}; // 起始坐标 (channel, height, width)
        std::vector<int> size = {y1_new - y0_new, x1_new - x0_new, 3};   // 裁剪大小 (channel, height, width)
        auto startVar = _Const(start.data(), {3}, NCHW, halide_type_of<int>());
        auto sizeVar = _Const(size.data(), {3}, NCHW, halide_type_of<int>());

        // 使用 _Slice 进行裁剪
        auto croppedImage = _Slice(original_image, startVar, sizeVar);
        // 输出裁剪后的图像信息
        auto croppedInfo = croppedImage->getInfo()->dim;
        int ih = croppedInfo[0];
        int iw = croppedInfo[1];
        int ch = croppedInfo[2];
        printf("Cropped image shape: %d x %d x %d\n", ih, iw, ch);
        int len = ih > iw ? ih : iw;
        float scale = len / 256;
        std::vector<int> padvals { 0, len - ih, 0, len - iw, 0, 0 };
        auto pads = _Const(static_cast<void*>(padvals.data()), {3, 2}, NCHW, halide_type_of<int>());
        auto image = _Pad(croppedImage, pads, CONSTANT);
        image = resize(image, Size(256, 256), 0, 0, INTER_LINEAR, -1, {0., 0., 0.}, {1./255., 1./255., 1./255.});
        auto input = _Unsqueeze(image, {0});
        input = _Convert(input, NC4HW4);
        auto outputs = net_keypoint->onForward({input});
        auto output = _Convert(outputs[0], NCHW);
        output = _Squeeze(output);
        int randomNumber = std::rand();
        std::string filename = "cropped_result_" + std::to_string(randomNumber) + ".jpg";//文件名字
        printf("dims = "); for (auto d : output->getInfo()->dim) printf("%d, ", d); printf("\n");
        for (int i = 0; i < 98; ++i){
            float x = output->readMap<float>()[i * 2];
            float y = output->readMap<float>()[i * 2 + 1];
            circle(croppedImage, {x * iw, y * ih}, 1, {0, 255, 0}, 2);
            imwrite(filename, croppedImage);
        }
        rectangle(original_image, {x0, y0}, {x1, y1}, {0, 0, 255}, 2);//画人脸框
    }
    if (imwrite("res.jpg", original_image)) {
        MNN_PRINT("result image write to `res.jpg`.\n");
    }
    rtmgr->updateCache();
    return 0;
}
